from __future__ import division
from __future__ import print_function

import numpy as np
import torch
from tqdm import tqdm

from tools.avgmeter import AverageMeter
from utils.eval_metrics import evaluate_reranking
from utils.utils import DeepSupervision


def trainer(model, trainloader, xent, tent, optimizer, use_gpu):
    model.train()
    xent_losses = AverageMeter()
    tent_losses = AverageMeter()
    losses = AverageMeter()

    for batch_idx, (imgs, pids, img_path) in enumerate(trainloader):
        optimizer.zero_grad()
        if use_gpu:
            imgs = imgs.cuda()
            pids = pids.cuda()
        outputs, features = model(imgs)

        if isinstance(outputs, (tuple, list)):
            xent_loss = DeepSupervision(xent, outputs, pids)
        else:
            xent_loss = xent(outputs, pids)

        if isinstance(features, (tuple, list)):
            tent_loss = DeepSupervision(tent, features, pids)
        else:
            tent_loss = tent(features, pids)

        xent_losses.update(xent_loss.item(), 1)
        tent_losses.update(tent_loss.item(), 1)

        loss = xent_loss + tent_loss
        loss.backward()
        optimizer.step()
        losses.update(loss.item(), 1)

        print("Batch {}/{}\t Loss {:.6f} ({:.6f}) xent Loss {:.6f} ({:.6f}), tent Loss {:.6f} ({:.6f})".format(
            batch_idx + 1, len(trainloader), losses.val, losses.avg, xent_losses.val, xent_losses.avg, tent_losses.val,
            tent_losses.avg))
    return losses.avg


def tester(model, queryloader, galleryloader, use_gpu, ranks=[1, 5, 10, 20], test_distance="cosine"):
    with torch.no_grad():
        model.eval()
        qf, q_pids = [], []
        for batch_idx, (imgs, pids, img_path) in enumerate(tqdm(queryloader)):
            del img_path
            if use_gpu:
                imgs = imgs.cuda()
                pids = pids.cuda()
            if len(imgs.size()) == 6:
                method = 'dense'
                b, n, s, c, h, w = imgs.size()
                assert (b == 1)
                imgs = imgs.view(b * n, s, c, h, w)
            else:
                method = None

            features = model(imgs)
            q_pids.extend(pids.data.cpu())

            features = features.data.cpu()
            torch.cuda.empty_cache()
            features = features.view(-1, features.size(1))

            if method == 'dense':
                features = torch.mean(features, 0, keepdim=True)
            qf.append(features)

        qf = torch.cat(qf, 0)
        q_pids = np.asarray(q_pids)

        print("Extracted features for query set, obtained {}-by-{} matrix".format(qf.size(0), qf.size(1)))

        gf, g_pids = [], []
        for batch_idx, (imgs, pids, img_path) in enumerate(tqdm(galleryloader)):
            if use_gpu:
                imgs = imgs.cuda()
                pids = pids.cuda()

            if len(imgs.size()) == 6:
                method = 'dense'
                b, n, s, c, h, w = imgs.size()
                assert (b == 1)
                imgs = imgs.view(b * n, s, c, h, w)
            else:
                method = None

            features = model(imgs)
            features = features.data.cpu()
            torch.cuda.empty_cache()
            features = features.view(-1, features.size(1))

            if method == 'dense':
                features = torch.mean(features, 0, keepdim=True)

            g_pids.extend(pids.data.cpu())
            gf.append(features)

        gf = torch.cat(gf, 0)
        g_pids = np.asarray(g_pids)

        print("Extracted features for gallery set, obtained {}-by-{} matrix".format(gf.size(0), gf.size(1)))
        print("Computing distance matrix")

        be_cmc, metrics = evaluate_reranking(qf, q_pids, gf, g_pids, ranks,
                                             test_distance)
        return metrics, be_cmc
